#include "dns.h"
#include "kozos.h"
#include "lib.h"
#include "udp.h"
#include "netdef.h"
#include "netlib.h"

typedef struct {
    uint16 transaction_id;
    uint16 flags;
    uint16 question_count;
    uint16 answer_count;
    uint16 authority_count;
    uint16 additional_count;
} dns_header;

static int dns_getIP(dns_header *dhdr, uint16 size, uint8 IP[], uint16 *tid)
{
    uint8 *p = (uint8 *)dhdr;
    uint16 i, ac;

    // headerをチェック
    if (!(size >= sizeof(*dhdr) &&
          (dhdr->flags & 0xf97f) == 0x8100 &&
          dhdr->question_count == 1 &&
          dhdr->answer_count > 0))
    {
        return -1;
    }

    i = sizeof(*dhdr); // pの添字。>= sizeはオーバーラン

    // question読み飛ばす
    while (1) {
        if (i >= size)
            return -1;
        uint8 len = p[i++];
        if (!len) break;
        i += len;
    }
    i += 4;

    // answerからIPを見つける
    for (ac = dhdr->answer_count; ac > 0; ac--) {
        uint16 type, class, dlen;

        while (1) {
            if (i >= size)
                return -1;
            if (p[i] < 64) {
                uint8 len = p[i++];
                if (!len) break;
                i += len;
            } else if (p[i] >= 192) { // 圧縮
                i += 2;
                break;
            } else {
                return -1;
            }
        }

        if (i + 14 > size)
            return -1;
        type = ((uint16)p[i] << 8) | p[i+1];
        i += 2;
        class = ((uint16)p[i] << 8) | p[i+1];
        i += 6;
        dlen = ((uint16)p[i] << 8) | p[i+1];
        i += 2;
        if (type == 1 && class == 1 && dlen == 4) {
            memcpy(IP, p+i, 4);
            if (tid) *tid = dhdr->transaction_id;
            return 0;
        } else {
            i += dlen;
        }
    }

    return -1;
}

static void dns_query(const char fqdn[], uint16 src_port, uint16 tid, const uint8 server[])
{
    const char *p;
    uint8 *q;
    dns_header *dhdr;
    int i, buflen, datalen;

    buflen = sizeof(dns_header) + strlen(fqdn) + 6;
    dhdr = kz_kmalloc(buflen);

    dhdr->transaction_id = tid;
    dhdr->flags = 0x0100;
    dhdr->question_count = 1;
    dhdr->answer_count = 0;
    dhdr->authority_count = 0;
    dhdr->additional_count = 0;

    p = fqdn;
    q = (uint8 *)(dhdr + 1);
    i = 0;
    while (1) {
        if (p[i] == '.') {
            *q++ = i;
            memcpy(q, p, i);
            q += i;
            p += i + 1;
            i = 0;
            continue;
        }

        if (!p[i]) {
            if (i > 0) {
                *q++ = i;
                memcpy(q, p, i);
                q += i;
            }
            *q++ = 0;
            break;
        }

        i++;
    }

    *q++ = 0; *q++ = 1; *q++ = 0; *q++ = 1;

    datalen = q - (uint8 *)dhdr;
    ASSERT(datalen == buflen || datalen == buflen - 1,
           "dns_query: Bad Datasize");
    udp_send((uint8 *)dhdr, datalen, src_port, UDP_PORT_DNS, server);

    kz_kmfree(dhdr);
}

int dns_resolve(const char fqdn[], uint8 IP[], int timeout, const uint8 server[])
{
    const udp_receiver *receiver;
    char *p;
    int ret = -1;

    if (!strcmp(fqdn, "localhost")) {
        IP[0] = 127; IP[1] = 0; IP[2] = 0; IP[3] = 1;
        return 0;
    }

    receiver = udp_use(0);
    if (!receiver) return -1;

    dns_query(fqdn, receiver->port, 0, server);

    if (kz_recva(receiver->mbox, NULL, &p, SECS(timeout)) != THREAD_ID_ERR) {
        ip_header *ihdr = (ip_header *)(p + sizeof(ethernet_header));
        udp_header *uhdr = (udp_header *)(ihdr + 1);
        if (!memcmp(ihdr->ip_src_IP, server, 4) && uhdr->udp_src_port == UDP_PORT_DNS)
            ret = dns_getIP((dns_header *)(uhdr + 1), uhdr->udp_length - sizeof(*uhdr), IP, NULL);
        kz_kmfree(p);
    }

    udp_close(receiver);

    return ret;
}
